from utils import get_input
import math
import matplotlib.pyplot as plt
import metric
import numpy as np


idx = 41
points_filenames = get_input.get_sampled_pointcloud(idx)
vibe_filenames = get_input.get_vibe_images(idx)
image_filenames = get_input.get_images(idx)

point_clouds = [get_input.read_point_cloud(points_filename)
                for points_filename in points_filenames]
nums = [len(points) for points in point_clouds]
distances = []
for points in point_clouds:
    centroid = np.mean(points, axis=0)
    x, _, z = centroid
    distances.append(int(math.sqrt(x ** 2 + z ** 2)))


pred_poses = get_input.get_pred_poses('pointnet2_stgcn', idx)
vibe_poses = get_input.get_vibe_poses(idx)
gt_poses = get_input.get_gt_poses(idx)


pred_poses = pred_poses[:len(gt_poses)]

pred_joints = get_input.poses_to_joints(pred_poses)
gt_joints = get_input.poses_to_joints(gt_poses)
mpjpes, _ = metric.compute_errors(gt_joints, pred_joints)
mpjpes = mpjpes.tolist()
frame_pcks = metric.compute_frame_pck(pred_joints, gt_joints, 0.3).tolist()

vibe_joints = get_input.poses_to_joints(vibe_poses)
vibe_mpjpes, _ = metric.compute_errors(gt_joints, vibe_joints)
vibe_mpjpes = vibe_mpjpes.tolist()
vibe_frame_pcks = metric.compute_frame_pck(
    vibe_joints, gt_joints, 0.3).tolist()

print(len(mpjpes))
print(len(vibe_mpjpes))


for i in range(len(mpjpes)):
    if distances[i] == 21 and mpjpes[i] > 0.048 and mpjpes[i] < 0.154 and vibe_mpjpes[i] > 0.162 and vibe_mpjpes[i] < 0.187:
        if frame_pcks[i] > 0.63 and frame_pcks[i] < 0.86 and vibe_frame_pcks[i] > 0.34 and vibe_frame_pcks[i] < 0.70:
            print(i, nums[i], vibe_filenames[i],
                  image_filenames[i], points_filenames[i])

# while True:
#     frame = int(input())
#     print(distances[frame], mpjpes[frame], frame_pcks[frame],
#           vibe_mpjpes[frame], vibe_frame_pcks[frame])

min_distances = min(distances)
max_distances = max(distances)


# PLOT
def get_metric_over_distance(metrics, distances):
    metrics_over_distance = [[] for _ in range((max_distances + 1))]
    for distance, m in zip(distances, metrics):
        metrics_over_distance[distance].append(m)

    res = []
    for distance_list in metrics_over_distance:
        if len(distance_list) == 0:
            res.append(0.0)
        else:
            res.append(np.mean(distance_list))

    return res[min_distances:]


mean_nums = get_metric_over_distance(nums, distances)
mean_mpjpes = [m * 1500 for m in get_metric_over_distance(mpjpes, distances)]
mean_frame_pcks = [
    m * 150 for m in get_metric_over_distance(frame_pcks, distances)]
mean_vibe_mpjpes = [
    m * 1500 for m in get_metric_over_distance(vibe_mpjpes, distances)]
mean_vibe_frame_pcks = [
    m * 150 for m in get_metric_over_distance(vibe_frame_pcks, distances)]


distances = list(range(min_distances, max_distances + 1))

plt.subplots(1, 3, figsize=(18, 3))
plt.subplot(1, 3, 1)
plt.xlabel('Distances from LiDAR(m)', fontsize=15)
plt.ylabel('Num of Points', fontsize=15)
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.plot(distances, mean_nums, color='darkseagreen', linewidth=4)
plt.subplot(1, 3, 2)
plt.xlabel('Distances from LiDAR(m)', fontsize=15)
plt.ylabel('MPJPE(mm)', fontsize=15)
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.plot(distances, mean_vibe_mpjpes, color='chocolate',
         label='VIBE', linestyle='--', linewidth=4)
plt.plot(distances, mean_mpjpes, color='steelblue',
         label='Ours', linewidth=4)
plt.legend(fontsize=15)
plt.subplot(1, 3, 3)
plt.xlabel('Distances from LiDAR(m)', fontsize=15)
plt.ylabel('PCK0.3(%)', fontsize=15)
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.plot(distances, mean_vibe_frame_pcks,
         color='orange', label='VIBE', linestyle='--', linewidth=4)
plt.plot(distances, mean_frame_pcks, color='lightcoral',
         label='Ours', linewidth=4)
plt.legend(fontsize=15)


plt.savefig('/home/ljl/tmp/dist_plot.png'.format(idx),
            dpi=600, bbox_inches='tight')
